#include <cube.h>

//determine row and column from blockIdx.x
CUBE_DEVICE(void, findPosition, int &Col, int &Row, int &blockX, int &blockY) {
    int k = blockIdx.x;
  blockY = -0.5 + sqrt(0.25 + 2*k);
  blockX = k - ((blockY+1)*(blockY))/2;
  Row = (blockY*TILE_HEIGHT + threadIdx.y);
  Col = (blockX*TILE_WIDTH + threadIdx.x);
}

CUBE_DEVICE(void, write1x1, int &Col, int &Row, cuFloatComplex *product, 
	    float sumXXreal, float sumXXimag, float sumXYreal, float sumXYimag,
	    float sumYXreal, float sumYXimag, float sumYYreal, float sumYYimag) {

  int f=blockIdx.y;

  if(Col>Row)return;
  
  ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 0] = 
    make_float4(sumXXreal, sumXXimag, sumXYreal, sumXYimag);
  ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 1] = 
    make_float4(sumYXreal, sumYXimag, sumYYreal, sumYYimag);
  
}

CUBE_DEVICE(void, write2x2, int &Col, int &Row, cuFloatComplex *product, 
	    float sum11XXreal, float sum11XXimag, float sum11XYreal, float sum11XYimag,
	    float sum11YXreal, float sum11YXimag, float sum11YYreal, float sum11YYimag,
	    float sum12XXreal, float sum12XXimag, float sum12XYreal, float sum12XYimag,
	    float sum12YXreal, float sum12YXimag, float sum12YYreal, float sum12YYimag,
	    float sum21XXreal, float sum21XXimag, float sum21XYreal, float sum21XYimag,
	    float sum21YXreal, float sum21YXimag, float sum21YYreal, float sum21YYimag,
	    float sum22XXreal, float sum22XXimag, float sum22XYreal, float sum22XYimag,
	    float sum22YXreal, float sum22YXimag, float sum22YYreal, float sum22YYimag){

  int f=blockIdx.y;

  if(Col>Row)return;

  ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 0] = 
    make_float4(sum11XXreal, sum11XXimag, sum11XYreal, sum11XYimag);
  ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 1] = 
    make_float4(sum11YXreal, sum11YXimag, sum11YYreal, sum11YYimag);
  ((float4*)product)[(f*NBASELINE + ((Row+1)*(Row+2)/2) + Col)*NPOL + 0] = 
    make_float4(sum21XXreal, sum21XXimag, sum21XYreal, sum21XYimag);
  ((float4*)product)[(f*NBASELINE + ((Row+1)*(Row+2)/2) + Col)*NPOL + 1] = 
    make_float4(sum21YXreal, sum21YXimag, sum21YYreal, sum21YYimag);
  ((float4*)product)[(f*NBASELINE + ((Row+1)*(Row+2)/2) + (Col+1))*NPOL + 0] = 
    make_float4(sum22XXreal, sum22XXimag, sum22XYreal, sum22XYimag);
  ((float4*)product)[(f*NBASELINE + ((Row+1)*(Row+2)/2) + (Col+1))*NPOL + 1] = 
    make_float4(sum22YXreal, sum22YXimag, sum22YYreal, sum22YYimag);
  
  // Test if entire tile needs to be written or just 3 of 4 parts (exclude top-right)
  if (Col<Row) {
    ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + (Col+1))*NPOL + 0] = 
      make_float4(sum12XXreal, sum12XXimag, sum12XYreal, sum12XYimag);
    ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + (Col+1))*NPOL + 1] = 
      make_float4(sum12YXreal, sum12YXimag, sum12YYreal, sum12YYimag);
  }

}

CUBE_DEVICE(void, write1x2, int &Col, int &Row, cuFloatComplex *product, 
	    float sum11XXreal, float sum11XXimag, float sum11XYreal, float sum11XYimag,
	    float sum11YXreal, float sum11YXimag, float sum11YYreal, float sum11YYimag,
	    float sum12XXreal, float sum12XXimag, float sum12XYreal, float sum12XYimag,
	    float sum12YXreal, float sum12YXimag, float sum12YYreal, float sum12YYimag){

  int f=blockIdx.y;

  if(Col>Row)return;

  ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 0] = 
    make_float4(sum11XXreal, sum11XXimag, sum11XYreal, sum11XYimag);
  ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 1] = 
    make_float4(sum11YXreal, sum11YXimag, sum11YYreal, sum11YYimag);
  
  // Test if entire tile needs to be written or just 1 of 2 parts (exclude right col)
  if (Col<Row) {
    ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + (Col+1))*NPOL + 0] = 
      make_float4(sum12XXreal, sum12XXimag, sum12XYreal, sum12XYimag);
    ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + (Col+1))*NPOL + 1] = 
      make_float4(sum12YXreal, sum12YXimag, sum12YYreal, sum12YYimag);
  }

}

CUBE_DEVICE(void, write2x1, int &Col, int &Row, cuFloatComplex *product, 
	    float sum11XXreal, float sum11XXimag, float sum11XYreal, float sum11XYimag,
	    float sum11YXreal, float sum11YXimag, float sum11YYreal, float sum11YYimag,
	    float sum21XXreal, float sum21XXimag, float sum21XYreal, float sum21XYimag,
	    float sum21YXreal, float sum21YXimag, float sum21YYreal, float sum21YYimag){

  int f=blockIdx.y;

  if(Col>Row+1)return;

  // Test if the top part should be written or just the bottom part
  if(Col<Row+1){
    ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 0] = 
      make_float4(sum11XXreal, sum11XXimag, sum11XYreal, sum11XYimag);
    ((float4*)product)[(f*NBASELINE + (Row*(Row+1)/2) + Col)*NPOL + 1] = 
      make_float4(sum11YXreal, sum11YXimag, sum11YYreal, sum11YYimag);
  }
  ((float4*)product)[(f*NBASELINE + ((Row+1)*(Row+2)/2) + Col)*NPOL + 0] = 
    make_float4(sum21XXreal, sum21XXimag, sum21XYreal, sum21XYimag);
  ((float4*)product)[(f*NBASELINE + ((Row+1)*(Row+2)/2) + Col)*NPOL + 1] = 
    make_float4(sum21YXreal, sum21YXimag, sum21YYreal, sum21YYimag);
}
